# @File    : ezmath.py

import math

import numpy
import scipy.interpolate

np = numpy
__EPS = 1e-9


def length_of_vec(v: numpy.ndarray):
    return numpy.sqrt((v * v).sum())


class Polator:
    def __init__(self, known_points, know_values):
        self.__extrapolator = scipy.interpolate.NearestNDInterpolator(known_points, know_values)
        self.__interpolator = None
        try:
            self.__interpolator = scipy.interpolate.LinearNDInterpolator(known_points, know_values)
        except scipy.spatial._qhull.QhullError as e:
            print("弃用内插，因为遇到异常：\n%s" % e)

    def __call__(self, *args, **kwargs):
        if self.__interpolator:
            res = self.__interpolator(*args, **kwargs)
            if numpy.isnan(res):
                res = self.__extrapolator(*args, )
        else:
            res = self.__extrapolator(*args, )
        return res


def RodriguesRotate(v: np.ndarray, u: np.ndarray, theta_rad: float) -> np.ndarray:
    u = u / np.linalg.norm(u)  # 计算单位向量
    sin_theta = np.sin(theta_rad)
    cos_theta = np.cos(theta_rad)
    v_new = v * cos_theta + np.cross(u, v) * sin_theta + np.dot(u, v) * u * (1 - cos_theta)
    return v_new


def gen_arc_from_p1_p2_rho(start_pt: numpy.ndarray, end_pt: numpy.ndarray, rho: float, d_s: float = 1,
                           n: numpy.ndarray = numpy.array([0, -1, 0])) -> tuple:
    mid_pt = (start_pt + end_pt) / 2.
    start_to_end = (end_pt - start_pt)
    a = length_of_vec(start_to_end) / 2.
    sin_half_theta = a / rho
    half_theta = math.asin(sin_half_theta)
    theta = half_theta * 2
    s_arc = rho * theta  # 总弧长
    d_theta = theta * d_s / s_arc
    _v_r = numpy.cross(start_to_end, n)
    _v_r_0 = _v_r / numpy.linalg.norm(_v_r)  # 圆心指向弧中点方向（r方向）上的单位向量
    center_pt = mid_pt - rho * math.cos(half_theta) * _v_r_0  # 圆心
    v_r_base = start_pt - center_pt
    res = []
    for theta_ in numpy.arange(0, theta, d_theta):
        res.append(RodriguesRotate(v_r_base, n, theta_) + center_pt)
    print("圆弧参数：偏转角度theta = %.2f rad，实际起点：%s，实际终点：%s，弧长：%s，转轴（圆弧所在平面的法线）：%s" % (theta, start_pt, end_pt, s_arc, n))
    return numpy.array(res), center_pt


def as_float_null_to_nan(arr: numpy.ndarray):
    arr_ = arr.copy()
    arr_[arr_ == ""] = numpy.nan
    return arr_.astype(float)


def fill_nan_for_perpendicular(v1: numpy.ndarray, v2_has_nan: numpy.ndarray):
    i = numpy.argwhere(numpy.isnan(v2_has_nan))[0][0]
    v1_mul_v2 = v1 * v2_has_nan
    temp = v1_mul_v2[:i].sum() + v1_mul_v2[i + 1:].sum()
    if abs(v1[i]) < __EPS:
        if abs(temp) < __EPS:
            v2_has_nan[i] = 100
            return
        else:
            raise RuntimeError("v1, v2不可能垂直")
    else:
        v2_has_nan[i] = -temp / v1[i]
    return
